import streamlit as st
from langchain_core.prompts import ChatPromptTemplate
from customize.get_ollama import GetOllama

prompt = ChatPromptTemplate.from_template("请写一篇关于{topic}的中文小故事，不超过200字")
model = GetOllama(ip=GetOllama.ailab_linux_ip, model_name="qwen2.5:14b", model_type=1)()
chain = prompt | model
st.sidebar.text_input("Hello world!")
with st.form('my_form'):
    text = st.text_area('输入主题关键词：', '小白兔')
    submitted = st.form_submit_button('提交')
    if submitted:
        st.info(chain.invoke({"topic": text}).content)
    chain.get_graph().print_ascii()
